import function_lib as fn
import tridiagonal_matrix_algorithm as tma
import math

#default values
a = 0; b = 1
N = 100
bcc = { 'a':[0,1,0], 'b':[1,0,0]}
#bcc = { 'a':[1,0,0], 'b':[1,-4,0]}

#Asking for a, b, N
reply1 = input('Set a,b and N manually instead of\
 using default a=0, b=1, N =100? [y]: ')

#manual set
if reply1 == 'y':
    a = float(input('Enter a: '))
    b = float(input('Enter b: '))
    N = int(input('Enter N: '))

#Asking for border conditions' coefficients
reply2 = input('Set border conditions\' coefficients\
 manually instead of using default (0,1,0) and (1,0,0)? [y]: ')

#manual set
if reply2 == 'y':
    bcc['a'][0] = input('Enter alpha1: ')
    bcc['a'][1] = input('Enter beta1: ')
    bcc['a'][2] = input('Enter gamma1: ')
    
    bcc['b'][0] = input('Enter alpha2: ')
    bcc['b'][1] = input('Enter beta2: ')
    bcc['b'][2] = input('Enter gamma2: ')

#Arrays
points = [a+i*(b-a)/N for i in range (N+1)]

p_val = [fn.p(points[i]) for i in range(N+1)]
q_val = [fn.q(points[i]) for i in range(N+1)]
f_val = [fn.f(points[i]) for i in range(N+1)]

#System
h = (b-a)/N

#make system_matrix
system_matrix = []

for i in range(N+1):
    #i - equation; j - point
    system_matrix.append([0 for j in range(N+1)])

#border equations
system_matrix[0][0] = bcc['a'][0] - bcc['a'][1]*(2-q_val[0]*h*h)\
/(h*(2-p_val[0]*h))
system_matrix[0][1] = 2*bcc['a'][1]/(h*(2-p_val[0]*h))

system_matrix[N][N-1] = -2*bcc['b'][1]/(h*(2+p_val[N]*h))
system_matrix[N][N] = bcc['b'][0] + bcc['b'][1]*(2-q_val[N]*h*h)\
/(h*(2+p_val[N]*h))

#middle
for i in range(1, N):
    system_matrix[i][i-1] = 1/(h*h) - p_val[i]/(2*h)
    system_matrix[i][i] = -2/(h*h) + q_val[i]
    system_matrix[i][i+1] = 1/(h*h) + p_val[i]/(2*h)

#f_vector
f_vector = [0 for i in range(N+1)]

f_vector[0] = bcc['a'][2] + bcc['a'][1]*f_val[0]*h/(2-p_val[0]*h)
f_vector[N] = bcc['b'][2] - bcc['b'][1]*f_val[N]*h/(2+p_val[N]*h)

for i in range(1, N):
    f_vector[i] = f_val[i]

#TMA
y_val = tma.TMA(system_matrix, f_vector)

#Output
f = open('data', 'w')

for i in range(N+1):
    f.write(str(points[i]) + ' ' + str(y_val[i]) + '\n')

f.close()
